{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eb231031",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c65a3dff",
   "metadata": {},
   "source": [
    "# Sparse Operation Kit with NVTabular Data Reader\n",
    "\n",
    "## Overview\n",
    "\n",
    "Sparse Operation Kit (hereafter SOK) is a toolkit aiming at wrapping sparse operation algorithms, commonly used in recommendation scenarios, into a user-friendly library. When users want to leverage those GPU-accelerated algorithms to speed up their application, they can quickly start from the Python toolkit.\n",
    "\n",
    "[NVTabular](https://github.com/NVIDIA/NVTabular) is a feature engineering and preprocessing library for tabular data that is designed to easily manipulate terabyte scale datasets and train deep learning (DL) based recommender systems. It's identified that the dataloader is one major bottleneck in deep learning recommender systems when being trained with TensorFlow. The dataloader cannot prepare the next batch fast enough and therefore, the GPU is not fully utilized.\n",
    "\n",
    "NVTabular dataloader’s features are:\n",
    "\n",
    "- removing bottleneck of item-by-item dataloading\n",
    "- enabling larger than memory dataset by streaming from disk\n",
    "- reading data directly into GPU memory and remove CPU-GPU communication\n",
    "- preparing batch asynchronously in GPU to avoid CPU-GPU communication\n",
    "- supporting commonly used .parquet format\n",
    "- easy integration into existing TensorFlow pipelines by using similar API - works with tf.keras models\n",
    "\n",
    "More information about NVTabular in [blogpost](https://medium.com/nvidia-merlin/training-deep-learning-based-recommender-systems-9x-faster-with-tensorflow-cc5a2572ea49).\n",
    "\n",
    "This notebook demonstrates how to use Sparse Operation Kit (SOK demo can be found [here](./sparse_operation_kit_demo.ipynb)) with NVTabular data reader. Instead of TensorFlow Dataset APIs, NVIDIA's NVTabular data loading library can alleviate existing bottlenecks and bring to bear the full power of GPU acceleration.\n",
    "\n",
    "### Model structure\n",
    "\n",
    "This demo model is constructed with a dense embedding layer and 7 fully connected layers, where the former 6 fully connected layers have 1024 outsput units, and the last one has 1 output unit.\n",
    "\n",
    "![model structure](https://raw.githubusercontent.com/NVIDIA/HugeCTR/master/sparse_operation_kit/documents/source/images/demo_model_structure.png)\n",
    "\n",
    "In our experiments, we see a speed-up by 3x on data-loading of the same training workflow with NVTabular dataloader for this SOK demo. \n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "- [Overview](#Overview)\n",
    "- [Installation](#Installation)\n",
    "- [Prepare training data](#Prepare-training-data)\n",
    "    - [Generate synthetic dataset](#Generate-synthetic-dataset)\n",
    "- [Convert data](#Convert-data)\n",
    "- [Training SOK model with NVTabular data reader](#Trainig-SOK-model-with-NVTabular-data-reader)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5071b6d6",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "### Requirements\n",
    "\n",
    "- TensorFlow == 2.x\n",
    "- NVTabular == 0.6.0\n",
    "\n",
    "### Get SOK from NGC\n",
    "\n",
    "The SparseOperationKit is preinstalled in the [Merlin Tensorflow Training Container](https://ngc.nvidia.com/catalog/containers/nvidia:merlin:merlin-tensorflow-training). You can simply start the notebook inside the running NGC docker container:\n",
    "`nvcr.io/nvidia/merlin/merlin-training:22.04`\n",
    "\n",
    "And you can check the existence of required libraries by running the following Python code after launching this container.\n",
    "\n",
    "```python\n",
    "$ python3 -c \"import sparse_operation_kit as sok\"\n",
    "```\n",
    "\n",
    "### Build SOK from Source Code\n",
    "\n",
    "Sometimes new features are not merged into NGC container in time, so you can build SOK from source code. If you want to build SparseOperationKit from the souce code instead of using the NGC container, please refer to the [Setup development environment](../../docs/hugectr_contributor_guide.md#build-sparse-operation-kit-sok-from-source-code)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43c31071",
   "metadata": {},
   "source": [
    "## Prepare training data\n",
    "\n",
    "We generate synthetic dataset which is more practical and provides better understanding of our SOK demo. It contains one dense feature and related label. We can set the number of feature fields in the used embedding layers and the number of keys in each slot.\n",
    "\n",
    "### Generate synthetic dataset\n",
    "\n",
    "First, specify hyper parameters for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7d5d5235",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f\n",
    "# Resets the namespace by removing all names defined by the user without asking for confirmation\n",
    "\n",
    "args = dict()\n",
    "\n",
    "args[\"gpu_num\"] = 8                               # the number of available GPUs\n",
    "args[\"global_batch_size\"] = 8192                  # the globally batchsize for all GPUs\n",
    "args[\"slot_num\"] = 100                            # the number of feature fields in this embedding layer\n",
    "args[\"nnz_per_slot\"] = 10                         # the number of keys in each slot\n",
    "args[\"vocabulary_size\"] = 1024 * 8\n",
    "args[\"iter_num\"] = 30                             # the number of training iterations\n",
    "args[\"demo_dirpath\"] = \"../documents/tutorials/DenseDemo/\" # the path to SOK demo\n",
    "args[\"data_dirpath\"] = args[\"demo_dirpath\"] + \"data/\"                           # the path used to save the generated data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a31ff8c",
   "metadata": {},
   "source": [
    "Then we will generate synthetic dataset and initial values that is used to initialize embedding parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84274398",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]: sparse_operation_kit is imported\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-09-13 07:08:53.058873: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:08:54.287132: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    }
   ],
   "source": [
    "# Standard Libraries\n",
    "import sys, os\n",
    "from pathlib import Path\n",
    "\n",
    "# External Dependencies\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# SOK python scripts\n",
    "sys.path.append(\"../\")\n",
    "sys.path.append(\"../unit_test/test_scripts/\")\n",
    "sys.path.append(\"../documents/tutorials/\")\n",
    "import utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8aec2893",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]: begin to generate random samples\n",
      "[INFO]: generated random samples\n",
      "[INFO]: begin to generate random samples\n",
      "[INFO]: generated random samples\n",
      "[INFO]: begin to generate random samples\n",
      "[INFO]: generated random samples\n",
      "[INFO]: total_samples.shape=(245760, 100, 10), total_labels.shape=(245760, 1)\n",
      "[INFO]: dumpped items to file ../documents/tutorials/DenseDemo/data/data.file\n"
     ]
    }
   ],
   "source": [
    "# Create directory for generated data\n",
    "Path(args[\"data_dirpath\"]).mkdir(parents=True, exist_ok=True)\n",
    "data_filename = args[\"data_dirpath\"] + \"data.file\"\n",
    "\n",
    "counts = args[\"iter_num\"] // 10\n",
    "total_samples, total_labels = None, None\n",
    "\n",
    "for _ in range(counts):\n",
    "    random_samples, random_labels = utility.generate_random_samples(\n",
    "                                        num_of_samples=args[\"global_batch_size\"] * 10,\n",
    "                                        vocabulary_size=args[\"vocabulary_size\"],\n",
    "                                        slot_num=args[\"slot_num\"],\n",
    "                                        max_nnz=args[\"nnz_per_slot\"],\n",
    "                                        use_sparse_mask=False)\n",
    "    if total_samples is None:\n",
    "        total_samples = random_samples\n",
    "        total_labels = random_labels\n",
    "    else:\n",
    "        total_samples = np.concatenate([total_samples, random_samples], axis=0)\n",
    "        total_labels = np.concatenate([total_labels, random_labels], axis=0)\n",
    "\n",
    "print(\"[INFO]: total_samples.shape={}, total_labels.shape={}\".format(total_samples.shape, total_labels.shape))\n",
    "utility.save_to_file(data_filename, total_samples, total_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36f3c852",
   "metadata": {},
   "source": [
    "## Convert data\n",
    "\n",
    "### Convert synthetic data to parquet file\n",
    "\n",
    "We are supposed to convert the data to parquet file first for use of NVTabular data reader. And we split the data into several parts for every worker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9df1734",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_labels = total_labels.reshape(-1) \n",
    "num_parts = args[\"gpu_num\"]\n",
    "samples = np.split(total_samples, num_parts, axis=0)\n",
    "labels = np.split(total_labels, num_parts, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a88e0c",
   "metadata": {},
   "source": [
    "We create several pandas DataFrames and then concatenante them for avoidance of limitation of Python list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0d5cec6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                 feature  label\n",
      "30715  [[77, 74, 68, 59, 55, 40, 29, 21, 12, 2], [159...      0\n",
      "30716  [[73, 72, 71, 46, 45, 44, 19, 19, 18, 17], [15...      0\n",
      "30717  [[69, 65, 53, 49, 37, 33, 29, 17, 12, 1], [154...      1\n",
      "30718  [[78, 78, 69, 69, 23, 23, 14, 14, 6, 5], [158,...      1\n",
      "30719  [[79, 72, 62, 62, 53, 35, 26, 17, 7, 0], [161,...      1\n",
      "                                                 feature  label\n",
      "30715  [[80, 80, 80, 58, 58, 42, 33, 21, 20, 13], [15...      1\n",
      "30716  [[79, 72, 68, 46, 43, 39, 35, 31, 6, 2], [151,...      1\n",
      "30717  [[66, 64, 63, 40, 38, 37, 35, 13, 11, 9], [154...      0\n",
      "30718  [[72, 70, 68, 67, 65, 64, 32, 29, 26, 25], [15...      1\n",
      "30719  [[76, 68, 67, 66, 53, 39, 24, 24, 9, 1], [136,...      0\n",
      "                                                 feature  label\n",
      "30715  [[77, 64, 64, 63, 61, 53, 25, 23, 22, 22], [16...      1\n",
      "30716  [[65, 47, 38, 30, 29, 20, 19, 12, 11, 2], [154...      0\n",
      "30717  [[80, 77, 62, 48, 47, 46, 31, 28, 15, 13], [16...      0\n",
      "30718  [[54, 53, 51, 50, 48, 15, 13, 12, 10, 8], [139...      1\n",
      "30719  [[78, 70, 64, 63, 48, 43, 34, 19, 12, 4], [154...      1\n",
      "                                                 feature  label\n",
      "30715  [[79, 67, 59, 47, 41, 39, 27, 18, 13, 5], [157...      0\n",
      "30716  [[80, 79, 68, 65, 54, 54, 39, 25, 13, 13], [16...      0\n",
      "30717  [[75, 66, 62, 46, 42, 33, 27, 18, 13, 9], [160...      1\n",
      "30718  [[80, 77, 72, 67, 26, 22, 17, 12, 9, 4], [154,...      0\n",
      "30719  [[78, 75, 65, 58, 54, 53, 43, 30, 18, 7], [148...      1\n",
      "                                                 feature  label\n",
      "30715  [[79, 73, 66, 62, 59, 56, 18, 10, 8, 5], [160,...      0\n",
      "30716  [[71, 65, 61, 55, 51, 35, 31, 21, 4, 0], [158,...      1\n",
      "30717  [[77, 75, 65, 56, 31, 21, 12, 10, 3, 0], [142,...      1\n",
      "30718  [[77, 77, 64, 58, 45, 32, 27, 27, 14, 9], [161...      0\n",
      "30719  [[77, 64, 59, 58, 51, 36, 31, 24, 23, 5], [155...      1\n",
      "                                                 feature  label\n",
      "30715  [[80, 70, 69, 66, 63, 60, 46, 36, 6, 2], [161,...      0\n",
      "30716  [[78, 69, 60, 51, 43, 42, 33, 24, 15, 15], [15...      0\n",
      "30717  [[72, 59, 46, 43, 33, 30, 20, 17, 7, 4], [152,...      0\n",
      "30718  [[69, 68, 68, 68, 54, 54, 53, 39, 39, 39], [16...      0\n",
      "30719  [[79, 29, 28, 27, 24, 20, 17, 13, 9, 2], [150,...      0\n",
      "                                                 feature  label\n",
      "30715  [[75, 70, 68, 63, 56, 51, 45, 13, 5, 1], [160,...      0\n",
      "30716  [[76, 57, 55, 48, 46, 36, 27, 25, 16, 6], [157...      0\n",
      "30717  [[71, 63, 55, 47, 41, 39, 33, 25, 13, 5], [161...      0\n",
      "30718  [[77, 72, 60, 52, 44, 41, 36, 32, 24, 7], [144...      0\n",
      "30719  [[80, 77, 75, 69, 63, 59, 53, 47, 5, 2], [153,...      1\n",
      "                                                 feature  label\n",
      "30715  [[69, 63, 56, 45, 38, 37, 31, 30, 13, 6], [158...      0\n",
      "30716  [[74, 72, 59, 56, 50, 48, 26, 24, 2, 0], [159,...      0\n",
      "30717  [[73, 72, 54, 51, 34, 33, 32, 13, 12, 11], [15...      1\n",
      "30718  [[76, 63, 53, 38, 29, 28, 19, 15, 6, 5], [159,...      0\n",
      "30719  [[72, 64, 64, 59, 56, 51, 48, 43, 40, 31], [15...      1\n"
     ]
    }
   ],
   "source": [
    "sample_name = \"feature\"\n",
    "label_name = \"label\"\n",
    "dfs = list()\n",
    "for i in range(num_parts):\n",
    "    dfs.append(pd.DataFrame({\n",
    "        sample_name : samples[i].tolist(),\n",
    "        label_name : labels[i].tolist(),\n",
    "    }))\n",
    "    print(dfs[-1].tail(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67c8b71f",
   "metadata": {},
   "source": [
    "We concatenante the DataFrames and the \"index\" column should be popped out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c8c157bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                  feature  label\n",
      "245755  [[69, 63, 56, 45, 38, 37, 31, 30, 13, 6], [158...      0\n",
      "245756  [[74, 72, 59, 56, 50, 48, 26, 24, 2, 0], [159,...      0\n",
      "245757  [[73, 72, 54, 51, 34, 33, 32, 13, 12, 11], [15...      1\n",
      "245758  [[76, 63, 53, 38, 29, 28, 19, 15, 6, 5], [159,...      0\n",
      "245759  [[72, 64, 64, 59, 56, 51, 48, 43, 40, 31], [15...      1\n"
     ]
    }
   ],
   "source": [
    "df = pd.concat(dfs)\n",
    "## Reset index\n",
    "df = df.reset_index()\n",
    "df.pop(\"index\")\n",
    "print(df.tail(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1309d243",
   "metadata": {},
   "source": [
    "We can simply transfer it to parquet file with Pandas API `to_parquet`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1d0b1d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "parquet_filename = args[\"data_dirpath\"] + \"data.parquet\"\n",
    "df.to_parquet(parquet_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "856ff88b",
   "metadata": {},
   "source": [
    "When MPI is used, we'd like to let each CPU process have its own datareader, and each datareader reads from different data source. Therefore the whole dataset is splited. We use `nvt.Workflow` to transform original data to split data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "38f6bb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split \n",
    "import nvtabular as nvt\n",
    "\n",
    "workflow = nvt.Workflow(nvt.ColumnGroup(sample_name) + label_name)\n",
    "workflow.transform(nvt.Dataset(parquet_filename)).to_parquet(\n",
    "    output_path=args[\"data_dirpath\"] + \"convert/\",\n",
    "    shuffle=None,\n",
    "    out_files_per_proc=num_parts,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "852db3a1",
   "metadata": {},
   "source": [
    "## Trainig SOK model with NVTabular data reader\n",
    "\n",
    "We train the SOK model with TensorFlow MultiWorkerMirroredStrategy and MPI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "edce78a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting run_sok_nvt.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile run_sok_nvt.py\n",
    "\n",
    "import argparse\n",
    "import sys,os,json,glob\n",
    "sys.path.append(\"../\")\n",
    "sys.path.append(\"../documents/tutorials/\")\n",
    "sys.path.append(\"../unit_test/test_scripts/\")\n",
    "sys.path.append(\"../documents/tutorials/DenseDemo/\")\n",
    "\n",
    "MPI_SIZE = int(os.getenv(\"OMPI_COMM_WORLD_SIZE\"))\n",
    "MPI_RANK = int(os.getenv(\"OMPI_COMM_WORLD_RANK\"))\n",
    "\n",
    "def main(args, task_id):\n",
    "    import nvtabular as nvt  # noqa: E402 isort:skip\n",
    "    from nvtabular.framework_utils.tensorflow import layers  # noqa: E402 isort:skip\n",
    "    from nvtabular.loader.tensorflow import KerasSequenceLoader  # noqa: E402 isort:skip\n",
    "    from nvtabular.io.dataset import Dataset\n",
    "\n",
    "    import tensorflow as tf\n",
    "    from models import SOKDenseDemo\n",
    "    import utility\n",
    "    from utility import sparse_operation_kit as sok\n",
    "    import nvtx\n",
    "\n",
    "\n",
    "    comm_options = tf.distribute.experimental.CommunicationOptions(\n",
    "        bytes_per_pack=0,\n",
    "        timeout_seconds=None,\n",
    "        implementation=tf.distribute.experimental.CommunicationImplementation.NCCL\n",
    "    )\n",
    "\n",
    "    if args.total_gpu_num == 1:\n",
    "        strategy = tf.distribute.MirroredStrategy()\n",
    "    else:\n",
    "        port = 12345\n",
    "        os.environ[\"TF_CONFIG\"] = json.dumps({\n",
    "            \"cluster\": {\"worker\": [\"localhost\" + \":\" + str(port + i) \n",
    "                                    for i in range(args.worker_num)]},\n",
    "            \"task\": {\"type\": \"worker\", \"index\": task_id}\n",
    "        })\n",
    "        strategy = tf.distribute.MultiWorkerMirroredStrategy(\n",
    "            communication_options=comm_options)\n",
    "\n",
    "    sample_name = \"feature\"\n",
    "    label_name = \"label\"\n",
    "\n",
    "    nvt_loader = KerasSequenceLoader(\n",
    "        paths_or_dataset=sorted(glob.glob(args.data_filename)),\n",
    "        batch_size=args.global_batch_size//MPI_SIZE,\n",
    "        label_names=[label_name],\n",
    "        cont_names=[sample_name],\n",
    "        engine=\"parquet\",\n",
    "        shuffle=False,\n",
    "        buffer_size=0.06,  # how many batches to load at once\n",
    "        sparse_as_dense=True,\n",
    "        global_size=MPI_SIZE,\n",
    "        global_rank=task_id,\n",
    "    )\n",
    "\n",
    "    with strategy.scope():\n",
    "        sok.Init(global_batch_size=args.global_batch_size)\n",
    "\n",
    "        model = SOKDenseDemo(max_vocabulary_size_per_gpu=args.max_vocabulary_size_per_gpu,\n",
    "                             embedding_vec_size=args.embedding_vec_size,\n",
    "                             slot_num=args.slot_num,\n",
    "                             nnz_per_slot=args.nnz_per_slot,\n",
    "                             num_dense_layers=args.num_dense_layers)\n",
    "\n",
    "        embedding_optimizer = utility.get_embedding_optimizer(args.optimizer)(learning_rate=0.1)\n",
    "        dense_optimizer = utility.get_dense_optimizer(args.optimizer)(learning_rate=0.1)\n",
    "\n",
    "    loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n",
    "    def _replica_loss(labels, logits):\n",
    "        loss = loss_fn(labels, logits)\n",
    "        return tf.nn.compute_average_loss(loss, global_batch_size=args.global_batch_size)\n",
    "\n",
    "    @tf.function\n",
    "    def _train_step(inputs, labels):\n",
    "        with tf.GradientTape() as tape:\n",
    "            logit = model(inputs, training=True)\n",
    "            loss = _replica_loss(labels, logit)\n",
    "        emb_variable, other_variable = sok.split_embedding_variable_from_others(model.trainable_variables)\n",
    "        grads, emb_grads = tape.gradient(loss, [other_variable, emb_variable])\n",
    "        if \"plugin\" not in args.optimizer:\n",
    "            with sok.OptimizerScope(emb_variable):\n",
    "                embedding_optimizer.apply_gradients(zip(emb_grads, emb_variable),\n",
    "                                                    experimental_aggregate_gradients=False)\n",
    "        else:\n",
    "            embedding_optimizer.apply_gradients(zip(emb_grads, emb_variable),\n",
    "                                                experimental_aggregate_gradients=False)\n",
    "        \n",
    "        # mannually all-reduce dense gradients\n",
    "        replica_context = tf.distribute.get_replica_context()\n",
    "        grads = replica_context.all_reduce(\"sum\", grads, \n",
    "                                            options=comm_options)\n",
    "        dense_optimizer.apply_gradients(zip(grads, other_variable),\n",
    "                                        experimental_aggregate_gradients=False)\n",
    "\n",
    "        # manually all-reduce loss, it is ok, because replica_loss has already been used to \n",
    "        # update local variables.\n",
    "        loss = replica_context.all_reduce(tf.distribute.ReduceOp.SUM, loss,\n",
    "                                          options=comm_options)\n",
    "        return loss\n",
    "\n",
    "    input_tensor_shape = [-1, args.slot_num, args.nnz_per_slot]\n",
    "    for i, (feats, labels) in enumerate(nvt_loader):\n",
    "        if args.stop_at_iter > 0 and i >= args.stop_at_iter:\n",
    "            break\n",
    "\n",
    "        rng = nvtx.start_range(message=\"Iteration_\" + str(i), color=\"blue\")\n",
    "        \n",
    "        input_vals, input_nnzs = feats[sample_name]\n",
    "        inputs = tf.reshape(input_vals, input_tensor_shape)\n",
    "        total_loss = strategy.run(_train_step, args=(inputs, labels))\n",
    "\n",
    "        nvtx.end_range(rng)\n",
    "        print(\"[INFO]: Iteration: {}, loss={}\".format(i, total_loss))\n",
    "\n",
    "# def set_affinity(rank):\n",
    "#     affinity_map = {0: list(range(48,64)) + list(range(176,192)),\n",
    "#                     1: list(range(48,64)) + list(range(176,192)),\n",
    "#                     2: list(range(16,32)) + list(range(144,160)),\n",
    "#                     3: list(range(16,32)) + list(range(144,160)),\n",
    "#                     4: list(range(112,128)) + list(range(240,256)),\n",
    "#                     5: list(range(112,128)) + list(range(240,256)),\n",
    "#                     6: list(range(80,96)) + list(range(208,224)),\n",
    "#                     7: list(range(80,96)) + list(range(208,224))}\n",
    "\n",
    "#     my_affinity = affinity_map[rank]\n",
    "#     import os\n",
    "#     os.sched_setaffinity(0, my_affinity)\n",
    "\n",
    "def set_affinity(rank):\n",
    "    num_part = MPI_SIZE\n",
    "    part_cpu_cnt = os.cpu_count() // num_part\n",
    "    total_affinity = os.sched_getaffinity(0)\n",
    "    affinity_map = dict()\n",
    "    for i in range(num_part):\n",
    "        tmp_map = list(total_affinity)[:part_cpu_cnt] if i % 2 == 0 else list(total_affinity)[part_cpu_cnt:]\n",
    "        affinity_map.update({i : tmp_map})\n",
    "\n",
    "    part = num_part // 2\n",
    "    my_affinity = affinity_map[rank % part] + affinity_map[(rank+part-1) % part + part]\n",
    "\n",
    "    try:\n",
    "        os.sched_setaffinity(0, my_affinity)\n",
    "    except OSError:\n",
    "        import traceback\n",
    "        print(\"[rank-{}] os.sched_setaffinity(0, {}) -> OSError\".format(rank, my_affinity))\n",
    "        traceback.print_exc()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser(description=\"run DNN model with SparseOperationKit\")\n",
    "\n",
    "    parser.add_argument(\"--data_filename\", type=str,\n",
    "                        help=\"the filename of training datas\",\n",
    "                        default=\"./data/convert/*.parquet\")\n",
    "    parser.add_argument(\"--global_batch_size\", type=int,\n",
    "                        required=True)\n",
    "    parser.add_argument(\"--max_vocabulary_size_per_gpu\", type=int,\n",
    "                        required=True)\n",
    "    parser.add_argument(\"--slot_num\", type=int, required=True,\n",
    "                        help=\"the number of feature fields\")\n",
    "    parser.add_argument(\"--nnz_per_slot\", type=int, required=True,\n",
    "                        help=\"the number of keys in each slot\")\n",
    "    parser.add_argument(\"--num_dense_layers\", type=int, required=True,\n",
    "                        help=\"the number of fully connected layers in this DNN model\")\n",
    "    parser.add_argument(\"--embedding_vec_size\", type=int, required=True,\n",
    "                        help=\"the dimension of embedding vectors\")\n",
    "    parser.add_argument('--optimizer', type=str,\n",
    "                        help=\"use what optimizer\",\n",
    "                        required=False, default='plugin_adam',\n",
    "                        choices=['plugin_adam', 'adam', 'sgd'])\n",
    "    parser.add_argument(\"--stop_at_iter\", type=int, required=False,\n",
    "                        help=\"early stop the process if iteration reachs this setting.\",\n",
    "                        default=-1)\n",
    "    parser.add_argument(\"--data_splited\", type=int, required=False,\n",
    "                        default=0, choices=[0, 1],\n",
    "                        help=\"it is a flag used to denotes whether the data is already splited.\"+\\\n",
    "                             \"by default, it is set to 0, which means the data is not splited.\")\n",
    "\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    size = MPI_SIZE\n",
    "    args.worker_num = size\n",
    "    args.total_gpu_num = size\n",
    "\n",
    "    task_id = MPI_RANK\n",
    "\n",
    "    set_affinity(task_id)\n",
    "\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(task_id)\n",
    "    main(args, task_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc291b6f",
   "metadata": {},
   "source": [
    "Add `--oversubscribe` to `mpiexec` if there is not enough slots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fdd5a86d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-09-13 07:10:51.231643: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.304766: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.304764: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.402127: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.402127: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.404622: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.429951: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:51.430007: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0\n",
      "2021-09-13 07:10:53.695597: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set\n",
      "2021-09-13 07:10:53.695763: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1\n",
      "2021-09-13 07:10:53.705794: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 0 with properties: \n",
      "pciBusID: 0000:89:00.0 name: Tesla V100-SXM2-32GB computeCapability: 7.0\n",
      "coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 31.75GiB deviceMemoryBandwidth: 836.37GiB/s\n",
      "...\n",
      "2021-09-13 07:11:04.517064: I sparse_operation_kit/kit_cc/kit_cc_infra/src/facade.cc:230] SparseOperationKit allocated internal memory.\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 0, loss=1.3331298828125\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 1, loss=976074.1875\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 2, loss=2216.47412109375\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "[INFO]: Iteration: 3, loss=0.7021139860153198\n",
      "...\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n",
      "[INFO]: Iteration: 29, loss=0.6933262944221497\n"
     ]
    }
   ],
   "source": [
    "!mpiexec -n 8 --allow-run-as-root \\\n",
    "    python3 run_sok_nvt.py \\\n",
    "    --data_filename=\"../documents/tutorials/DenseDemo/data/convert/*.parquet\" \\\n",
    "    --global_batch_size=8192 \\\n",
    "    --max_vocabulary_size_per_gpu=8192 \\\n",
    "    --slot_num=100 \\\n",
    "    --nnz_per_slot=10 \\\n",
    "    --num_dense_layers=6 \\\n",
    "    --embedding_vec_size=4 \\\n",
    "    --data_splited=1 \\\n",
    "    --optimizer=\"adam\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f7760f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
